import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 1. 数据准备
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
stock_list = ['600519.SH', '000858.SZ', '601318.SH', '600036.SH', '000333.SZ']
start_date = '20210101'
end_date = '20230101'


# 获取数据
def get_data():
    stock_data = []
    for stock in stock_list:
        df = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
        df['代码'] = stock
        stock_data.append(df)
    return pd.concat(stock_data)


all_data = get_data()
all_data['trade_date'] = pd.to_datetime(all_data['trade_date'])
all_data = all_data.rename(columns={
    'trade_date': '日期', 'close': '收盘', 'high': '最高',
    'low': '最低', 'vol': '成交量'
})


# 2. 特征工程与标签创建
def create_features(df):
    # 未来5日收益率标签
    df['未来5日收盘'] = df.groupby('代码')['收盘'].shift(-5)
    df['未来5日收益率'] = (df['未来5日收盘'] - df['收盘']) / df['收盘']
    df['标签'] = np.where(df['未来5日收益率'] > 0, 1, 0)

    # 技术指标
    df['5日均线'] = df['收盘'].rolling(5).mean()
    df['20日均线'] = df['收盘'].rolling(20).mean()
    df['RSI'] = 100 - (100 / (1 + (df['收盘'].diff(1).clip(lower=0).rolling(14).mean() /
                                   -df['收盘'].diff(1).clip(upper=0).rolling(14).mean())))
    df['MACD'] = df['收盘'].ewm(span=12).mean() - df['收盘'].ewm(span=26).mean()
    df['波动率'] = df['收盘'].pct_change().rolling(20).std()
    df['成交量_20日均值'] = df['成交量'].rolling(20).mean()
    return df


all_data = all_data.groupby('代码').apply(create_features).dropna()

# 3. 数据预处理与可视化
numeric_cols = ['收盘', '5日均线', '20日均线', 'RSI', 'MACD', '波动率', '成交量_20日均值']

# 异常值处理可视化
plt.figure(figsize=(12, 6))
sns.boxplot(data=all_data[numeric_cols])
plt.title('技术指标异常值检测(处理前)')
plt.xticks(rotation=45)
plt.savefig('tech_indicators_boxplot.png', bbox_inches='tight')  # 保存图像
plt.show()

# 异常值处理
for col in numeric_cols:
    q1 = all_data[col].quantile(0.25)
    q3 = all_data[col].quantile(0.75)
    iqr = q3 - q1
    all_data = all_data[(all_data[col] > q1 - 1.5 * iqr) & (all_data[col] < q3 + 1.5 * iqr)]

# 归一化
scaler = MinMaxScaler()
all_data[numeric_cols] = scaler.fit_transform(all_data[numeric_cols])

# 主成分分析
pca = PCA(n_components=0.95)
pca_features = pca.fit_transform(all_data[numeric_cols])

# PCA可视化
plt.figure(figsize=(10, 6))
plt.bar(range(1, pca.n_components_ + 1), pca.explained_variance_ratio_)
plt.xlabel('主成分')
plt.ylabel('解释方差比例')
plt.title('PCA分析 - 主成分解释方差比例')
plt.savefig('pca_explained_variance.png', bbox_inches='tight')  # 保存图像
plt.show()

# 相关性分析
plt.figure(figsize=(10, 8))
sns.heatmap(all_data[numeric_cols].corr(), annot=True, cmap='coolwarm', center=0)
plt.title('技术指标相关性矩阵')
plt.savefig('correlation_heatmap.png', bbox_inches='tight')  # 保存图像
plt.show()

# 个股画像可视化
sample_stock = np.random.choice(stock_list)
profile_cols = ['RSI', 'MACD', '波动率', '成交量_20日均值', '20日均线', '收盘']
stock_profile = all_data[all_data['代码'] == sample_stock].copy()

# 确保日期是datetime格式并排序
stock_profile['日期'] = pd.to_datetime(stock_profile['日期'])
stock_profile = stock_profile.sort_values('日期')

plt.figure(figsize=(16, 10))  # 增大图形尺寸

for i, col in enumerate(profile_cols):
    plt.subplot(2, 3, i + 1)
    plt.plot(stock_profile['日期'], stock_profile[col], marker='o', markersize=3, linewidth=1)  # 添加标记和调整线宽
    plt.title(f'{col}走势')
    plt.xticks(rotation=45)

    # 调整y轴范围，留出一些边距
    y_min, y_max = stock_profile[col].min(), stock_profile[col].max()
    plt.ylim(y_min - 0.1 * (y_max - y_min), y_max + 0.1 * (y_max - y_min))

    # 添加网格线
    plt.grid(True, linestyle='--', alpha=0.5)

plt.tight_layout()
plt.suptitle(f'股票{sample_stock}技术指标画像', y=1.02, fontsize=14)
plt.savefig(f'{sample_stock}_profile.png', bbox_inches='tight', dpi=300)  # 提高保存分辨率
plt.show()

# 4. 数据均衡与建模
X = all_data[numeric_cols]
y = all_data['标签']

# 数据分布可视化
plt.figure(figsize=(6, 4))
y.value_counts().plot(kind='bar')
plt.title('标签分布 (1:上涨, 0:下跌)')
plt.show()

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 随机森林建模
rf = RandomForestClassifier(n_estimators=100, random_state=42, class_weight="balanced")
rf.fit(X_train, y_train)
y_pred = rf.predict(X_test)
y_prob = rf.predict_proba(X_test)[:, 1]

# 5. 模型评价可视化
# 分类报告
print(classification_report(y_test, y_pred))

# 混淆矩阵
plt.figure(figsize=(6, 6))
sns.heatmap(confusion_matrix(y_test, y_pred), annot=True, fmt='d', cmap='Blues')
plt.title('混淆矩阵')
plt.xlabel('预测标签')
plt.ylabel('真实标签')
plt.savefig('confusion_matrix.png', bbox_inches='tight')  # 保存图像
plt.show()

# ROC曲线
fpr, tpr, _ = roc_curve(y_test, y_prob)
roc_auc = auc(fpr, tpr)

plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.2f})')
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('假阳性率')
plt.ylabel('真阳性率')
plt.title('ROC曲线')
plt.legend(loc="lower right")
plt.savefig('roc_curve.png', bbox_inches='tight')  # 保存图像
plt.show()

# 特征重要性
feature_imp = pd.DataFrame({
    '特征': numeric_cols,
    '重要性': rf.feature_importances_
}).sort_values('重要性', ascending=False)

plt.figure(figsize=(10, 6))
sns.barplot(x='重要性', y='特征', data=feature_imp)
plt.title('特征重要性排序')
plt.savefig('feature_importance.png', bbox_inches='tight')  # 保存图像
plt.show()